import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsClassifier
from sklearn.ensemble import RandomForestClassifier

f = open("wine.data")
l1 = f.readlines()
f.close()
characteristics = np.zeros((len(l1), 13))  # 数据集的属性，每个样本有13种属性
wineType = np.zeros((len(l1)))  # 红酒的种类，1,2,3三种。
for i, r in enumerate(l1):
    l2 = r.split(',')
    wineType[i] = int(l2[0])
    for j, strn in enumerate(l2[1:]):
        characteristics[i, j] = float(strn)
# 以上的步骤是生成数据集
print('数据集的数量', len(characteristics))  # 13个元素，分别代表酒精度等。wine.data一共十四列，以逗号分隔。
# wine.data文件的第一列是酒的类别代号，后面13列是酒精度、气味、沉淀物等指标。
print('红酒的类型', wineType)
X_train, X_test, y_train, y_test = train_test_split(characteristics, wineType, test_size=0.2)  # 将文件分为训练集和测试集。
print("训练集：", X_train, y_train)

knn = KNeighborsClassifier(algorithm='ball_tree')  # 近邻分类器

knn.fit(X_train, y_train)  # 使用K-近邻分类器

params = knn.get_params()
score = knn.score(X_test, y_test)
print("K-近邻分类器预测得分为：%s" % score)
print(knn.predict(X_test))  # k-近邻分类后的结果
print(y_test)  # 正确的结果

rfc = RandomForestClassifier()  # 使用随机森林分类器
rfc.fit(X_train, y_train)
params = rfc.get_params()
score2 = rfc.score(X_test, y_test)
print("随机森林预测得分为：%s" % score2)
print(rfc.predict(X_test))  # 随机森林分类后的结果
print(y_test)  # 正确的结果

import joblib  # 存储模型

joblib.dump(rfc, '红酒分类.pkl')
